from typing import Dict, List, Optional, Tuple
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from abc import ABC, abstractmethod
import ccxt
import time

class DataHandler(ABC):
    """数据处理器基类"""
    
    @abstractmethod
    def get_latest_bar(self, exchange: str, symbol: str) -> Optional[Dict]:
        """获取最新的K线数据"""
        pass
        
    @abstractmethod
    def get_latest_bars(self, exchange: str, symbol: str, N: int = 1) -> List[Dict]:
        """获取最近N根K线数据"""
        pass
        
    @abstractmethod
    def update_data(self):
        """更新数据"""
        pass
        
    @abstractmethod
    def has_more_data(self) -> bool:
        """是否还有更多数据"""
        pass

class HistoricalDataHandler(DataHandler):
    """历史数据处理器"""
    
    def __init__(self, exchanges: Dict[str, ccxt.Exchange], symbols: List[str],
                 start_date: datetime, end_date: datetime, timeframe: str = '1h'):
        """初始化历史数据处理器"""
        self.exchanges = exchanges
        self.symbols = symbols
        self.start_date = start_date
        self.end_date = end_date
        self.timeframe = timeframe
        
        self.data = {}  # 存储所有数据
        self.current_idx = {}  # 当前数据索引
        self.current_date = start_date
        
        self._load_data()
        
    def _fetch_ohlcv_with_retry(self, exchange: ccxt.Exchange, symbol: str, 
                               since: int, limit: int = None, max_retries: int = 3) -> List:
        """带重试的K线数据获取"""
        for attempt in range(max_retries):
            try:
                ohlcv = exchange.fetch_ohlcv(
                    symbol,
                    timeframe=self.timeframe,
                    since=since,
                    limit=limit
                )
                if ohlcv and len(ohlcv) > 0:
                    return ohlcv
                time.sleep(1)  # 避免频繁请求
            except Exception as e:
                print(f"获取K线数据失败 (尝试 {attempt + 1}/{max_retries}): {str(e)}")
                if attempt < max_retries - 1:
                    time.sleep(2 ** attempt)  # 指数退避
                else:
                    raise
        return []
        
    def _load_data(self):
        """加载历史数据"""
        print("\n加载历史数据...")
        for exchange_name, exchange in self.exchanges.items():
            for symbol in self.symbols:
                print(f"正在从{exchange_name}获取{symbol}的历史数据...")
                
                # 计算时间范围
                since = int(self.start_date.timestamp() * 1000)
                end_ts = int(self.end_date.timestamp() * 1000)
                
                all_ohlcv = []
                while since < end_ts:
                    # 获取一批数据
                    ohlcv = self._fetch_ohlcv_with_retry(exchange, symbol, since)
                    if not ohlcv:
                        break
                        
                    all_ohlcv.extend(ohlcv)
                    
                    # 更新since时间戳
                    last_timestamp = ohlcv[-1][0]
                    if last_timestamp <= since:
                        break
                    since = last_timestamp
                    
                if not all_ohlcv:
                    print(f"警告: 未能获取到{symbol}的历史数据")
                    continue
                    
                # 转换为DataFrame
                df = pd.DataFrame(
                    all_ohlcv,
                    columns=['timestamp', 'open', 'high', 'low', 'close', 'volume']
                )
                
                # 转换时间戳并设置索引
                df['timestamp'] = pd.to_datetime(df['timestamp'], unit='ms')
                df.set_index('timestamp', inplace=True)
                
                # 去重和排序
                df = df[~df.index.duplicated(keep='first')]
                df.sort_index(inplace=True)
                
                # 过滤时间范围
                mask = (df.index >= self.start_date) & (df.index <= self.end_date)
                df = df.loc[mask]
                
                print(f"获取到{len(df)}条{symbol}的历史数据")
                
                # 存储数据
                self.data[(exchange_name, symbol)] = df
                self.current_idx[(exchange_name, symbol)] = 0
                
    def get_latest_bar(self, exchange: str, symbol: str) -> Optional[Dict]:
        """获取最新的K线数据"""
        key = (exchange, symbol)
        if key not in self.data:
            return None
            
        idx = self.current_idx[key]
        if idx >= len(self.data[key]):
            return None
            
        row = self.data[key].iloc[idx]
        return {
            'timestamp': row.name,
            'open': row['open'],
            'high': row['high'],
            'low': row['low'],
            'close': row['close'],
            'volume': row['volume']
        }
        
    def get_latest_bars(self, exchange: str, symbol: str, N: int = 1) -> List[Dict]:
        """
        获取最近N根K线数据
        
        Args:
            exchange: 交易所名称
            symbol: 交易对
            N: 获取数据条数
            
        Returns:
            最近N根K线数据
        """
        key = (exchange, symbol)
        if key not in self.data:
            return []
            
        idx = self.current_idx[key]
        if idx >= len(self.data[key]):
            return []
            
        # 获取最近N根K线数据
        start_idx = max(0, idx - N + 1)
        bars = self.data[key].iloc[start_idx:idx + 1]
        
        # 转换为字典列表格式
        bars_list = []
        for timestamp, row in bars.iterrows():
            bar_dict = {
                'timestamp': timestamp,
                'open': row['open'],
                'high': row['high'],
                'low': row['low'],
                'close': row['close'],
                'volume': row['volume']
            }
            bars_list.append(bar_dict)
            
        return bars_list
        
    def update_data(self):
        """更新数据，模拟数据前进一个时间单位"""
        # 更新当前时间
        if self.timeframe == '1m':
            self.current_date += timedelta(minutes=1)
        elif self.timeframe == '5m':
            self.current_date += timedelta(minutes=5)
        elif self.timeframe == '15m':
            self.current_date += timedelta(minutes=15)
        elif self.timeframe == '1h':
            self.current_date += timedelta(hours=1)
        elif self.timeframe == '4h':
            self.current_date += timedelta(hours=4)
        elif self.timeframe == '1d':
            self.current_date += timedelta(days=1)
            
        # 更新所有数据的索引
        for key in self.current_idx:
            df = self.data[key]
            while (self.current_idx[key] < len(df) and 
                   df.index[self.current_idx[key]] <= self.current_date):
                self.current_idx[key] += 1
                
    def has_more_data(self) -> bool:
        """检查是否还有更多数据"""
        return any(idx < len(self.data[key]) 
                  for key, idx in self.current_idx.items())
                  
class LiveDataHandler(DataHandler):
    """实时数据处理器"""
    
    def __init__(self, exchanges: Dict[str, ccxt.Exchange], symbols: List[str],
                 timeframe: str = '1h'):
        """
        初始化实时数据处理器
        
        Args:
            exchanges: 交易所实例字典
            symbols: 交易对列表
            timeframe: K线周期
        """
        self.exchanges = exchanges
        self.symbols = symbols
        self.timeframe = timeframe
        self.data = {}
        
    def get_latest_bar(self, exchange: str, symbol: str) -> Optional[Dict]:
        """获取最新的K线数据"""
        if (exchange, symbol) not in self.data:
            return None
        return self.data[(exchange, symbol)][-1]
        
    def get_latest_bars(self, exchange: str, symbol: str, N: int = 1) -> pd.DataFrame:
        """获取最近N根K线数据"""
        if (exchange, symbol) not in self.data:
            return pd.DataFrame()
        return pd.DataFrame(self.data[(exchange, symbol)][-N:])
        
    def update_data(self):
        """更新实时数据"""
        for exchange_name, exchange in self.exchanges.items():
            for symbol in self.symbols:
                try:
                    # 获取最新K线数据
                    ohlcv = exchange.fetch_ohlcv(
                        symbol,
                        timeframe=self.timeframe,
                        limit=1
                    )
                    
                    if ohlcv:
                        bar = {
                            'timestamp': pd.to_datetime(ohlcv[0][0], unit='ms'),
                            'open': ohlcv[0][1],
                            'high': ohlcv[0][2],
                            'low': ohlcv[0][3],
                            'close': ohlcv[0][4],
                            'volume': ohlcv[0][5]
                        }
                        
                        key = (exchange_name, symbol)
                        if key not in self.data:
                            self.data[key] = []
                        self.data[key].append(bar)
                        
                except Exception as e:
                    print(f"Error updating data for {exchange_name} {symbol}: {e}")
                    
    def has_more_data(self) -> bool:
        """实时数据总是有效"""
        return True
